import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv, SAGEConv, GATConv
import torch.nn as nn
class GAT(torch.nn.Module):
    def __init__(self, node_size,features, hidden, classes, heads):
        super(GAT, self).__init__()
        self.gat1 = GATConv(features, hidden, heads=heads)
        self.gat2 = GATConv(hidden*heads, classes)
        self.fc1=nn.Linear(10,1)
        self.fc2 = nn.Linear(node_size, 1)

    def forward(self, x,edge_index):
        x = self.gat1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.gat2(x, edge_index)
        x=self.fc1(x)
        x=torch.squeeze(x,1)
        x=self.fc2(x)
        return x
